/* -*- mode: c; tab-width: 4; indent-tabs-mode: nil -*- */

/*
 * spawn: Run a batch of N shell commands in parallel.
 * Copyright (C) 2007  Simo Melenius <sme@iki.fi>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
#define _POSIX2_SOURCE
#define _XOPEN_SOURCE 500
#include <stdio.h>
#include <stdlib.h>
#include <getopt.h>
#include <assert.h>
#include <signal.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>
#include <errno.h>
#include <string.h>

#define USAGE "Usage: %s [-h] [-v] [-jN]"
#define MIN_INPUT_BUFFER_SIZE 128
#define MAX(a, b) (((a) > (b)) ? (a) : (b))

/*
 * Runtime data and configuration.
 */

int         numRunning  = 0;
int         exitCode    = EXIT_SUCCESS;
int         verbose     = 0;

/*
 * Function prototypes.
 */

int                     main            (int argc, char* const * argv);
static int              die             (const char* errormsg);
static void             handle_signal   (int sig);
static int              runJobs         (int maxJobs);
static char*            readLine        (void);

/*
 * Spawn.
 */

int main (int argc, char* const * argv)
{
    int maxJobs = 0;
    int opt;

    /* Read arguments. */

    while (-1 != (opt = getopt(argc, argv, "j::vh")))
    {
        switch (opt)
        {
        case 'j':
            if (optarg)
                maxJobs = MAX(0, atoi(optarg));
            break;
        case 'v':
            verbose = 1;
            break;
        case 'h':
            printf("\n" USAGE "\n\n", argv[0]);
            printf("Parallelize shell commands:\n\n");
            printf("  Read command lines from stdin, one per line, execute them as child processes\n");
            printf("  in parallel, and wait for all children to exit. If a child process fails\n");
            printf("  with a non-zero exit code, stop spawning new children, wait for all existing\n");
            printf("  children to exit and return the failed exit code.\n\n");
            printf("  -h    Display this help\n");
            printf("  -v    Print each command line to stdout after execution.\n");
            printf("  -jN   Spawn N parallel processes, 0 for unlimited. Defaults to 0.\n\n");
            printf("Version " VERSION "\n");
            exit(EXIT_SUCCESS);
        default:
            fprintf(stderr, "*** " USAGE "\n", argv[0]);
            exit(EXIT_FAILURE);
        }
    }

    /* Setup C-c trap. */

    signal(SIGINT, handle_signal);

    /* Run queue. */

    return runJobs(maxJobs);
}

/*
 * Emergency exit.
 */

int die (const char* errormsg)
{
    perror(errormsg);
    exit(exitCode ? exitCode : EXIT_FAILURE);
}

/*
 * Trap the first C-c for gracious exiting.
 */

void handle_signal (int sig)
{
    if (numRunning > 0)
    {
        fprintf(stderr, "*** terminating: waiting for children to die\n");
        fflush(stderr);
    }
    exitCode = EXIT_FAILURE;
}

/*
 * Read commands from stdin, fork at most ``maxJobs'' subprocesses at a
 * time to execute them, and wait for all to finish before returning.
 */

int runJobs (int maxJobs)
{
    char* line;
    int status;

    /* Main execution loop. */

    for (line = readLine(); exitCode == EXIT_SUCCESS && line; line = readLine())
    {
        int pid;

        /* If we have too many jobs running, wait for a child to exit. */

        if (maxJobs > 0)
        {
            while (numRunning == maxJobs)
            {
                if (wait(&status) != -1)
                {
                    numRunning--;

                    if (WEXITSTATUS(status) != EXIT_SUCCESS && exitCode == EXIT_SUCCESS)
                        exitCode = WEXITSTATUS(status);
                }
            }
        }

        /* In case of SIGINT or error, stop spawning children and wait to
         * reap the live ones; otherwise continue into forking. */

        if (exitCode == EXIT_SUCCESS)
        {
            /* Initialize and fork child. */

            pid = fork();

            switch (pid)
            {
            case 0: /* => Child. */
                status = system(line);

                /* Print command line in verbose mode: useful for teeing. */
                if (verbose)
                    fprintf(stdout, "%s\n", line); fflush(stdout);
                free(line);

                return (status == -1) ? 1 : WEXITSTATUS(status);

            case -1: /* => Failure. */
                die("fork() failed");
                break;

            default: /* => Success. */
                numRunning++;
            }
        }

        free(line);
    }

    /* Clean up and wait for existing processes. */

    while (numRunning > 0)
    {
        if (wait(&status) != 1)
        {
            if (WEXITSTATUS(status) != EXIT_SUCCESS && exitCode == EXIT_SUCCESS)
                exitCode = WEXITSTATUS(status);

            numRunning--;
        }
        else if (errno == ECHILD)
        {
            break;
        }
    }

    return exitCode;
}

/*
 * Return an malloc()'d buffer containing one line from stdin, without the
 * trailing newline.
 */

char* readLine (void)
{
    int     size = 0;
    int     maxsize = 0;
    char*   buffer = NULL;
    int     c;

    do
    {
        /* Read character, exit on EOF. */
        c = fgetc(stdin);
        if (c == EOF)
            break;

        /* Grow buffer in powers of two, if needed. */
        if (size == maxsize)
        {
            char* newbuf;
            maxsize = MAX(MIN_INPUT_BUFFER_SIZE, maxsize << 1);
            newbuf = realloc(buffer, maxsize);

            if (newbuf || die("out of memory"))
                buffer = newbuf;
        }

        /* Write non-linefeed characters into buffer. */
        if (c != '\n')
        {
            assert(size < maxsize);
            buffer[size++] = c;
        }
    } while (c != '\n');

    /* Null-terminate. */
    if (buffer)
        buffer[size] = '\0';

    return buffer;
}
